{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llm_rag.chains import build_pdf_retriever, Reranker, build_rag_chain\n",
    "from llm_rag.document import PDFSplitter\n",
    "from llm_rag.llms import ChatGLM\n",
    "import llm_rag.templates as templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = build_pdf_retriever()\n",
    "reranker = Reranker(reranker_dir='plms/bge-reranker-base/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6eb701a4d3754c219e7905a5f0f30cec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "llm = ChatGLM()\n",
    "llm.load_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "rag = build_rag_chain(retriever=retriever, reranker=reranker, llm=llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "RAG_TEMPLATE= \"\"\"根据汽车手册中的相关知识回复用户的问题.\n",
    "手册内容:\n",
    "{context}\n",
    "问题:\n",
    "{question}\n",
    "要求:\n",
    "1. 答案要包含所有正确关键词.\n",
    "2. 答案要简洁.\n",
    "3. 答案要尽量多的包含原文,不要自己编造不在原文中的内容.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "genereta (19 tokens, 0.59 sec, 32.00 tokens/sec)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'应该在车辆行驶一段时间后，当车漆失去光泽时为车身打蜡。'"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q = '什么情况应该为车身打蜡'\n",
    "rag.with_config(configurable={'rag_template': templates.RAG_TEMPLATE_C}).invoke(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='控灯功能后，打开车门时点亮阅读灯，关闭所有车门一段时间\\n后，阅读灯熄灭。\\n在中央显示屏中点击\\n-座舱体验-快捷设置，进入快捷设置界面。\\n仪表和灯光\\n89\\n01 点击开启/关闭门控灯。\\n使用超车灯\\n超车时，您可将拨杆朝向自身方向拉动然后松开，此时远光灯闪烁一\\n次，以提醒前方车辆。\\n使用危险警告灯\\n危险警告灯按键\\n车辆遇到交通事故或其他紧急情况时，按下危险警告灯按键，启用危\\n险警告灯。\\n说明！\\n□ 发生碰撞情况下，危险警告灯也将自动点亮。\\n仪表和灯光\\n90\\n调节背光亮度\\n使用调光旋钮调节亮度\\n白天情况下，未开启整车背光联动时，转动调光旋钮仅调节开关背光\\n亮度，开', metadata={'num_tokens': 256}),\n",
       " Document(page_content='盲点监测故障警告灯：盲点监测系统出现故障时，该警告\\n灯常亮。盲点监测系统标定未完成时，该警告灯闪烁。\\n低速行人提示音功能故障警告灯：低速行人提示音功能故\\n障时，该警告灯点亮。\\n指示灯和警告灯\\n指示灯图标\\n信息/通知用黄色显示：此信息指示灯点亮，同时仪表显\\n示屏上出现对应文本信息。该指示灯也可能会与其他图标\\n一同点亮。\\nREADY指示灯：车辆准备就绪时，该指示灯点亮。\\n转向指示灯：拨动转向指示灯开关时，对应的\\n指示灯闪烁。打开危险警告灯开关时，两侧转\\n向指示灯均闪烁。\\n近光灯指示灯：打开近光灯时，该指示灯点亮。\\n远光灯指示灯：打开远光灯时，该', metadata={'num_tokens': 256}),\n",
       " Document(page_content='反应。\\n注意！\\n■ 只有Lynk & Co领克中心能对CMSF系统进行维修，否则可能会导\\n致系统无法正常使用。\\n后方碰撞预警系统\\n当车辆行驶速度处于 0 - 150 km/h范围时，若系统检测到车辆存在后\\n方碰撞风险，将通过声音和视觉信号提醒驾驶员，并开启危险警告灯\\n以提醒后方车辆驾驶员注意制动。\\n当车辆处于静止状态，后方碰撞预警系统检测到无法避免后方碰撞时\\n将对本车进行制动，从而避免出现二次碰撞降低人员伤亡风险。\\n01 声音警告信号\\n02可视警告信号\\n驾驶辅助\\n213\\n警告！\\n■ 后方碰撞预警系统对驾驶员来说是一种辅助功能，您始终有责任\\n与其他车辆保持适当距离。\\n■ 后方碰撞预警系', metadata={'num_tokens': 256}),\n",
       " Document(page_content='动力电池电量。\\n06能量回收等级：显示车辆当前能量回收等级（自动、低、中、\\n高）。\\n07驾驶模式：显示当前车辆的驾驶模式。\\n08燃油量表：显示油箱中剩余的燃油油位。\\n09剩余燃油可续航里程：显示剩余燃油可行驶的距离。\\n指示灯和警告灯\\n指示灯图标\\n信息/通知用黄色显示：此信息指示灯点亮，同时仪表显\\n示屏上出现对应文本信息。该指示灯也可能会与其他图标\\n一同点亮。\\nREADY指示灯：车辆准备就绪时，该指示灯点亮。\\n转向指示灯：拨动转向指示灯开关时，对应的\\n指示灯闪烁。打开危险警告灯开关时，两侧转\\n向指示灯均闪烁。\\n近光灯指示灯：打开近光灯时，该指示灯点亮。', metadata={'num_tokens': 256}),\n",
       " Document(page_content='的\\n指示灯闪烁。打开危险警告灯开关时，两侧转\\n向指示灯均闪烁。\\n近光灯指示灯：打开近光灯时，该指示灯点亮。\\n远光灯指示灯：打开远光灯时，该指示灯点亮。\\n仪表和灯光\\n70\\n智能远近光控制开启指示灯：打开智能远近光控制时，该\\n指示灯点亮。\\n大灯随动转向功能激活指示灯：大灯随动转向功能激活\\n时，该指示灯点亮。\\n后雾灯指示灯：打开后雾灯时，该指示灯点亮。\\n位置灯指示灯：打开位置灯时，该指示灯点亮。\\n泊车紧急制动指示灯：泊车紧急制动激活时，该指示灯点\\n亮。\\n车道辅助系统启用指示灯：车道辅助系统启用时，该指示\\n灯点亮。\\n低速行人提示音功能关闭指示灯：低速', metadata={'num_tokens': 256})]"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs = retriever.get_relevant_documents(q)\n",
    "reranker(inputs={'input': q, 'docs': docs})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
